package com.wurd.youdao;

import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;

import java.net.InetSocketAddress;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;
import java.util.function.Supplier;

public class YoudaoSocket extends WebSocketServer {

    private ConcurrentHashMap<Integer, Translator> sockets = new ConcurrentHashMap<>();

    private int capacity;

    private ArrayBlockingQueue<Object[]> queue;

    private final Object wait = new Object();

    public void setSelector(Supplier<Translator> selector) {
        this.selector = selector;
    }

    private Supplier<Translator> selector;

    private Translator select() {
        synchronized (wait) {
            Translator translator = selector.get();
            if (translator == null) {
                try {
                    wait.wait();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                translator = selector.get();
            }
            return translator;
        }
    }

    public YoudaoSocket(int port) {
        this(new InetSocketAddress(port), 100);
    }

    public YoudaoSocket(int port, int capacity) {
        this(new InetSocketAddress(port), capacity);
    }

    public YoudaoSocket(InetSocketAddress address, int capacity) {
        super(address);
        this.capacity = capacity;
        this.queue = new ArrayBlockingQueue(this.capacity);
        this.selector = () -> {
            for (Translator translator : sockets.values()) {
                if (!translator.inUse && !translator.webSocket.isClosed()) {
                    translator.inUse = true;
                    return translator;
                }
            }
            return null;
        };
    }

    public void get(String str, Consumer<String> consumer) throws Exception {
        if (sockets.isEmpty()) {
            throw new Exception("socket is null");
        }
        // TODO translate using a thread pool
        Translator t = select();
        if (t != null) {
            t.data = str;
            t.consumer = consumer;
            t.translate();
        }
    }

    public boolean translate(String str, Consumer<String> consumer) throws InterruptedException {
        Object[] d = new Object[2];
        d[0] = str;
        d[1] = consumer;
        return translate(d);
    }

    public boolean translate(Object[] d) throws InterruptedException {
        return queue.offer(d);
    }

    public static class Translator {
        private String data;
        private boolean inUse;
        private final WebSocket webSocket;
        private Consumer<String> consumer;

        Translator(WebSocket webSocket) {
            this.webSocket = webSocket;
        }

        public void translate() {
            try {
                this.webSocket.send(data);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {
        System.out.println("open: " + webSocket.getRemoteSocketAddress().getHostString());
        synchronized (wait) {
            sockets.put(webSocket.hashCode(), new Translator(webSocket));
            wait.notifyAll();
        }
    }

    @Override
    public void onMessage(WebSocket webSocket, String s) {
        Translator translator = sockets.get(webSocket.hashCode());
        if (translator.consumer != null) {
            translator.consumer.accept(s);
        }
        translator.inUse = false;
        synchronized (wait) {
            wait.notifyAll();
        }
    }

    @Override
    public void onClose(WebSocket webSocket, int i, String s, boolean b) {
        System.out.println("close: " + webSocket.hashCode());
        sockets.remove(webSocket.hashCode());
    }

    @Override
    public void onError(WebSocket webSocket, Exception e) {
        System.out.println("error: " + webSocket.hashCode());
        e.printStackTrace();
        sockets.remove(webSocket.hashCode());
    }

    @Override
    public void onStart() {
        System.out.println("start...");
        genWorker();
    }

    private void genWorker() {
        new Thread(() -> {
            while (true) {
                try {
                    Object[] d = queue.take();
                    tran((String) d[0], (Consumer<String>) d[1]);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }).start();
    }

    private void tran(String d, Consumer<String> consumer) throws InterruptedException {
        Translator t = select();
        t.data = d;
        t.consumer = consumer;
        t.translate();
    }
}
